import torch
from PIL import Image

from transformers import DeepseekOcrForConditionalGeneration, DeepseekOcrProcessor
from transformers import model_addition_debugger_context


processor = DeepseekOcrProcessor.from_pretrained("deepseek_ocr_converted")
model = DeepseekOcrForConditionalGeneration.from_pretrained("deepseek_ocr_converted", dtype=torch.bfloat16)

image = Image.open("handwritten_letter_small.png").convert("RGB")

conversation = [
    {
        "role": "<|User|>",
        "content": [
            {"type": "image", "path": "./handwritten_letter_small.png"},
            {"type": "text", "text": "<|grounding|>Convert the document to markdown."},
        ],
    }
]

inputs = processor.apply_chat_template(
    conversation,
    return_dict=True,
    tokenize=True,
    add_generation_prompt=True,
    return_tensors="pt"
    )

with torch.no_grad():
    generated = model.generate(**inputs, max_new_tokens=50)

text = processor.batch_decode(generated, skip_special_tokens=False)[0]
print(text.strip())